import torch

class BasalModel(torch.nn.Module):
    def __init__(self):
        super(BasalModel, self).__init__()
        self.conv1 = torch.nn.Sequential(
            torch.nn.Conv2d(1, 64, kernel_size=3, stride=1, padding=1),
            torch.nn.ReLU())
        self.conv2 = torch.nn.Sequential(
            torch.nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
            torch.nn.ReLU())
        self.pool = torch.nn.MaxPool2d(stride=2, kernel_size=2)
        self.dense = torch.nn.Sequential(torch.nn.Linear(7 * 7 * 128, 1024),
                                         torch.nn.ReLU(),
                                         torch.nn.Dropout(p=0.5),
                                         torch.nn.Linear(1024, 200),
                                         torch.nn.ReLU(),
                                         torch.nn.Dropout(p=0.5),
                                         torch.nn.Linear(200, 10),
                                         torch.nn.Softmax(1))

    def forward(self, x):
        x = self.conv1(x)
        x = self.pool(x)
        x = self.conv2(x)
        x = self.pool(x)
        x = self.dense(x.view(-1, 7 * 7 * 128))
        return x
